#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
洪水预报智能代理

基于LangChain Agent框架实现的智能对话代理，整合参数提取和验证功能
"""

from langchain.agents import AgentExecutor, create_react_agent
from langchain.prompts import PromptTemplate
from langchain.agents.output_parsers import ReActSingleInputOutputParser
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.schema import AgentAction, AgentFinish
from langchain.memory import ConversationBufferMemory
import re
# from langchain_openai import ChatOpenAI  # 已删除OpenAI支持
from langchain_community.llms import Ollama
from langchain_community.chat_models import ChatOllama
from langchain.schema import HumanMessage, AIMessage, SystemMessage
from typing import Dict, Any, List, Optional, Tuple, Union
import json
from loguru import logger
from datetime import datetime


# 导入自定义工具
from parameter_validator_tool import ParameterValidatorTool
from parameter_update_tool import ParameterUpdateTool

# 导入配置
try:
    from config.config import OLLAMA_CONFIG
except ImportError:
    # 如果配置文件不存在，使用默认配置
    OLLAMA_CONFIG = {
        "base_url": "http://10.48.0.81:11434",
        "embedding_model": "bge-m3:latest",
        "chat_model": "qwq:latest"
    }

class QwqReActOutputParser(ReActSingleInputOutputParser):
    """自定义输出解析器，兼容qwq模型的<think>标签"""
    
    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
        """解析Agent输出，处理<think>标签"""
        # 移除<think>标签及其内容，但保留其他部分
        cleaned_text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
        
        # 调用父类的解析方法
        try:
            return super().parse(cleaned_text)
        except Exception as e:
            # 如果解析失败，尝试更激进的清理
            logger.warning(f"首次解析失败，尝试更激进的清理: {e}")
            
            # 提取最后的Action和Action Input
            action_match = re.search(r'Action:\s*([^\n]+)', cleaned_text)
            action_input_match = re.search(r'Action Input:\s*([^\n]+)', cleaned_text)
            
            if action_match and action_input_match:
                action = action_match.group(1).strip()
                action_input = action_input_match.group(1).strip()
                
                # 尝试解析action_input为JSON
                try:
                    import json
                    if action_input.startswith('{') and action_input.endswith('}'):
                        action_input_dict = json.loads(action_input)
                    else:
                        action_input_dict = action_input
                except:
                    action_input_dict = action_input
                
                return AgentAction(tool=action, tool_input=action_input_dict, log=text)
            
            # 检查是否是Final Answer
            final_answer_match = re.search(r'Final Answer:\s*(.+)', cleaned_text, re.DOTALL)
            if final_answer_match:
                return AgentFinish(return_values={"output": final_answer_match.group(1).strip()}, log=text)
            
            # 如果都失败了，重新抛出异常
            raise e

class FloodForecastAgent:
    """洪水预报智能代理
    
    整合参数提取、验证和对话管理功能的智能代理
    """
    
    def __init__(self, 
                 llm_config: Dict[str, Any],
                 max_iterations: int = None,
                 max_execution_time: float = None):
        """
        初始化智能代理
        
        Args:
            llm_config: LLM配置信息
            max_iterations: 最大迭代次数
            max_execution_time: 最大执行时间（秒）
        """
        self.llm_config = llm_config
        self.max_iterations = max_iterations
        self.max_execution_time = max_execution_time
        
        # 初始化工具
        self.validator_tool = ParameterValidatorTool()
        self.parameter_update_tool = ParameterUpdateTool(agent_instance=self)
        self.tools = [self.validator_tool, self.parameter_update_tool]
        
        # 初始化LLM
        self.llm = self._init_llm()
        
        # 状态管理
        self.collected_parameters = {}
        self.chat_history = []
        self.current_stage = "collection"  # collection, validation, confirmation, completed
        self.validation_result = None
        self.is_session_completed = False  # 标记会话是否已完成
        
        # 初始化记忆功能
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            input_key="input",
            output_key="output"
        )
        
        # 延迟初始化agent executor
        self.agent_executor = None
    
    def _init_llm(self):
        """初始化LLM - 仅支持Ollama"""
        return ChatOllama(
            model=self.llm_config.get('model', OLLAMA_CONFIG['chat_model']),
            temperature=self.llm_config.get('temperature', 0.1),
            base_url=self.llm_config.get('base_url', OLLAMA_CONFIG['base_url'])
        )
    
    def _create_system_prompt(self) -> str:
        """动态创建系统提示词"""
        # 基础系统提示
        base_prompt = """
你是一个专业的洪水预报参数收集助手。你的任务是通过自然对话收集用户的洪水预报需求参数。

必需的参数包括：
1. reservoir_name (水库名称): 需要进行洪水预报的水库名称
2. forecast_model (预报模型): 使用的预报模型，如新安江模型、萨克拉门托模型等
3. start_time (开始时间): 预报开始时间，格式为YYYY-MM-DD HH:MM
4. forecast_duration (预报时长): 预报持续时间，单位为小时
5. rainfall_forecast_period (降雨预见期): 降雨预报的时间长度，单位为小时
6. target_variable (目标变量): 预报的目标变量，如流量、水位等

工作流程：
1. 参数收集阶段：直接从用户输入中理解和提取参数信息
2. 参数验证阶段：所有参数收集完成后，调用ParameterValidatorTool验证参数
3. 参数确认阶段：与用户确认参数是否正确
4. 完成阶段：参数确认无误后，可以自由对话

参数提取能力：
- 你具备强大的自然语言理解能力，能直接从用户输入中识别和提取参数
- 能理解各种表达方式，如"预报24小时"、"从明天开始"、"大伙房水库"等
- 能处理时间格式转换，如"2025年8月7号下午13点"转为"2025-08-07 13:00"
- 能识别数值参数，如预报时长、降雨预见期等

工具使用的核心原则（渐进式收集策略）：
1. 【最高优先级】当用户提供明确的参数信息时，立即调用update_parameters工具更新，支持渐进式收集
2. 【次高优先级】对于模糊或不完整的信息，使用Final Answer反问澄清具体细节
3. 【基础原则】在参数验证阶段，调用ParameterValidatorTool验证收集到的参数
4. 【渐进原则】更新参数后继续询问剩余缺失参数，支持一点一点完善信息

参数更新工具的使用规则：
- 【立即更新原则】当用户提供明确的参数信息时，立即调用update_parameters工具更新
- 【允许条件】以下情况可以调用update_parameters工具：
  * "大伙房水库" → reservoir_name
  * "新安江模型" → forecast_model  
  * "2025-08-07 08:00" → start_time
  * "预报时长48小时" → forecast_duration
  * "降雨预见期12小时" → rainfall_forecast_period
  * "流量" → target_variable
- 【禁止条件】以下情况严禁调用update_parameters工具：
  * 时间信息不完整（如只说"明天"而未指定具体时分）
  * 数值信息模糊（如只说"24小时"而未明确是哪个参数）
  * 任何需要推测或补全的信息
- 【渐进式收集】更新参数后，继续询问剩余的缺失参数，支持一点一点收集信息

ReAct格式规则：
- 需要调用工具时：Thought: [简洁思考] → Action: [工具名称] → Action Input: [输入内容]
- 直接回复用户时：Thought: [简洁思考] → Final Answer: [回复内容]
- 思考过程要简洁明了，避免冗长的分析
- 禁止使用 "Action: 无" 或 "Action: 直接提问" 等无效格式
- 每个Thought后必须跟有效的Action或Final Answer
- 不要在思考中使用<think>标签或过度详细的分析

对话指导原则：
1. 使用自然、友好的语言与用户交流
2. 可以一次询问多个参数，但必须明确区分每个参数
3. 对用户输入进行确认和澄清
4. 提供参数的合理范围和示例
5. 在收集完所有参数后进行总结确认

参数收集规则：
- 明确参数立即更新，模糊信息反问澄清
- 支持渐进式收集，逐步完善参数
- 每次询问1-2个参数，避免信息过载
"""
        
        # 根据当前阶段和已收集参数动态调整提示
        if self.current_stage == "collection":
            missing_params = self._get_missing_parameters()
            base_prompt += "\n\n当前阶段：参数收集"
            base_prompt += "\n\n处理逻辑：明确参数立即更新，模糊信息反问澄清，逐步收集缺失参数。"
            
            if missing_params:
                base_prompt += f"\n\n当前缺少的参数: {', '.join(missing_params)}"
                base_prompt += "\n\n询问策略：每次询问1-2个参数，提供格式示例。"
            
            if self.collected_parameters:
                # 过滤掉内部字段，避免在模板中产生变量冲突
                filtered_params = {k: v for k, v in self.collected_parameters.items() 
                                 if k not in ['extraction_confidence', 'raw_input']}
                params_json = json.dumps(filtered_params, ensure_ascii=False)
                # 转义花括号以避免模板变量冲突
                params_json = params_json.replace('{', '{{').replace('}', '}}')
                base_prompt += f"\n已收集的参数: {params_json}"
        
        elif self.current_stage == "validation":
            base_prompt += "\n\n当前阶段：参数验证"
            base_prompt += "\n请调用ParameterValidatorTool验证所有参数的有效性"
            base_prompt += "\n重要：调用时需要将参数包装在 'parameters' 字段中，格式如：{{\"parameters\": {{参数字典}}}}"
            base_prompt += "\n验证完成后，根据验证结果与用户确认参数或要求修正"
            # 过滤掉内部字段，避免在模板中产生变量冲突
            filtered_params = {k: v for k, v in self.collected_parameters.items() 
                             if k not in ['extraction_confidence', 'raw_input']}
            params_json = json.dumps(filtered_params, ensure_ascii=False)
            # 转义花括号以避免模板变量冲突
            params_json = params_json.replace('{', '{{').replace('}', '}}')
            base_prompt += f"\n需要验证的参数: {params_json}"
        
        elif self.current_stage == "confirmation":
            base_prompt += "\n\n当前阶段：参数确认"
            base_prompt += "\n请向用户确认所有参数是否正确"
            base_prompt += "\n如果用户确认无误，则完成参数收集"
            base_prompt += "\n如果用户要求修改，请记录修改内容并重新进入收集阶段"
            base_prompt += "\n在此阶段可以直接与用户对话，无需强制调用工具"
            # 过滤掉内部字段，避免在模板中产生变量冲突
            filtered_params = {k: v for k, v in self.collected_parameters.items() 
                             if k not in ['extraction_confidence', 'raw_input']}
            params_json = json.dumps(filtered_params, ensure_ascii=False)
            # 转义花括号以避免模板变量冲突
            params_json = params_json.replace('{', '{{').replace('}', '}}')
            base_prompt += f"\n最终参数: {params_json}"
        
        elif self.current_stage == "completed":
            base_prompt += "\n\n当前阶段：已完成"
            base_prompt += "\n参数收集已完成！"
            base_prompt += "\n可以与用户进行自由对话，回答问题或提供帮助"
            base_prompt += "\n如果用户需要重新开始，请提示他们可以重置会话"
            # 过滤掉内部字段，避免在模板中产生变量冲突
            filtered_params = {k: v for k, v in self.collected_parameters.items() 
                             if k not in ['extraction_confidence', 'raw_input']}
            params_json = json.dumps(filtered_params, ensure_ascii=False)
            # 转义花括号以避免模板变量冲突
            params_json = params_json.replace('{', '{{').replace('}', '}}')
            base_prompt += f"\n已完成的参数: {params_json}"
        
        return base_prompt
    
    def _get_missing_parameters(self) -> List[str]:
        """获取缺失的参数"""
        required_params = [
            'reservoir_name', 'forecast_model', 'start_time',
            'forecast_duration', 'rainfall_forecast_period', 'target_variable'
        ]
        
        return [param for param in required_params 
                if param not in self.collected_parameters or 
                self.collected_parameters[param] is None]
    
    def _init_agent_executor(self) -> AgentExecutor:
        """初始化Agent执行器 - 仅支持Ollama ReAct代理"""
        # 创建动态提示模板
        system_prompt = self._create_system_prompt()
        
        # 创建ReAct提示模板，兼容qwq模型的<think>标签，并集成记忆功能
        template = f"""{system_prompt}

你有以下工具可以使用：
{{tools}}

对话历史：
{{chat_history}}

使用以下格式（可以使用<think>标签进行思考）：

Question: 用户的输入问题
<think>
你的详细思考过程，请参考对话历史来理解上下文和之前的交互
</think>

Thought: 简洁总结你要做什么
Action: 要采取的行动，应该是[{{tool_names}}]中的一个
Action Input: 行动的输入
Observation: 行动的结果
... (这个Thought/Action/Action Input/Observation可以重复N次)
<think>
最终思考，结合对话历史和当前状态
</think>

Thought: 简洁总结
Final Answer: 对原始输入问题的最终答案

重要：
1. 可以在<think>标签中进行详细思考
2. Thought部分要简洁明了
3. Action必须是工具列表中的一个
4. Action Input必须是有效的JSON格式
5. 请充分利用对话历史来理解用户意图和当前进度

开始！

Question: {{input}}
{{agent_scratchpad}}"""
        
        prompt = PromptTemplate.from_template(template)
        agent = create_react_agent(self.llm, self.tools, prompt)
        
        # 创建包装的LLM来预处理输出
        class QwqLLMWrapper:
            def __init__(self, llm):
                self.llm = llm
                
            def __getattr__(self, name):
                return getattr(self.llm, name)
                
            def invoke(self, input_data, config=None, **kwargs):
                """调用LLM并预处理输出"""
                result = self.llm.invoke(input_data, config, **kwargs)
                
                # 如果结果是字符串，清理<think>标签
                if isinstance(result, str):
                    cleaned_result = self._clean_think_tags(result)
                    return cleaned_result
                # 如果结果有content属性，清理content
                elif hasattr(result, 'content'):
                    result.content = self._clean_think_tags(result.content)
                    return result
                else:
                    return result
                    
            def _clean_think_tags(self, text: str) -> str:
                """智能清理<think>标签，保持ReAct格式完整性"""
                # 移除<think>标签及其内容
                cleaned = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
                
                # 清理多余的空行和空格
                lines = cleaned.split('\n')
                cleaned_lines = []
                for line in lines:
                    stripped = line.strip()
                    if stripped:  # 只保留非空行
                        cleaned_lines.append(stripped)
                
                # 重新组合，确保格式正确
                result = '\n'.join(cleaned_lines)
                
                # 确保Action和Action Input在正确的行上
                result = re.sub(r'\n+Action:', '\nAction:', result)
                result = re.sub(r'\n+Action Input:', '\nAction Input:', result)
                
                return result
        
        # 使用包装的LLM重新创建agent
        wrapped_llm = QwqLLMWrapper(self.llm)
        agent = create_react_agent(wrapped_llm, self.tools, prompt)
        
        # 创建agent executor，增强错误处理并集成记忆功能
        return AgentExecutor(
            agent=agent,
            tools=self.tools,
            memory=self.memory,
            max_iterations=self.max_iterations,
            max_execution_time=self.max_execution_time,
            verbose=True,
            return_intermediate_steps=True,
            handle_parsing_errors="Check your output and make sure it conforms to the format instructions. If you use <think> tags, make sure they don't interfere with Action and Action Input sections."
        )
    
    def process_user_input(self, user_input: str) -> Dict[str, Any]:
        """处理用户输入 - 统一LLM调用架构"""
        try:
            logger.info(f"处理用户输入: {user_input}")
            
            # 重新初始化agent executor以使用最新的系统提示
            self.agent_executor = self._init_agent_executor()
            
            # 对于Ollama ReAct代理，传入input，记忆会自动管理对话历史
            result = self.agent_executor.invoke({
                "input": user_input
            })
            
            # 同步更新内部chat_history以保持兼容性
            self.chat_history.append(HumanMessage(content=user_input))
            self.chat_history.append(AIMessage(content=result["output"]))
            
            # 处理工具执行结果
            self._process_tool_results(result.get("intermediate_steps", []))
            
            # 更新阶段状态
            self._update_stage()
            
            # 构建响应
            response = {
                "response": result["output"],
                "collected_parameters": self.collected_parameters,
                "missing_parameters": self._get_missing_parameters(),
                "current_stage": self.current_stage,
                "validation_result": self.validation_result,
                "is_complete": self._check_completion(),
                "confidence_score": self._calculate_confidence_score(),
                "timestamp": datetime.now().isoformat()
            }
            
            logger.info(f"处理完成: {response}")
            return response
            
        except Exception as e:
            logger.error(f"处理用户输入失败: {str(e)}")
            return {
                "response": f"抱歉，处理您的请求时出现错误: {str(e)}",
                "collected_parameters": self.collected_parameters,
                "missing_parameters": self._get_missing_parameters(),
                "current_stage": self.current_stage,
                "validation_result": None,
                "is_complete": False,
                "confidence_score": 0.0,
                "timestamp": datetime.now().isoformat(),
                "error": str(e)
            }
    
    def update_parameters_directly(self, user_input: str) -> Dict[str, Any]:
        """保留的旧方法，现在通过工具调用实现参数更新
        
        注意：此方法已被新的统一架构替代，参数更新现在通过update_parameters工具完成
        """
        logger.warning("update_parameters_directly方法已被弃用，请使用统一的Agent调用架构")
        return {}
    
    def _process_tool_results(self, intermediate_steps: List[Tuple]):
        """处理工具执行结果"""
        for action, observation in intermediate_steps:
            tool_name = action.tool
            
            if tool_name == "ParameterValidatorTool":
                try:
                    self.validation_result = json.loads(observation)
                    logger.info(f"参数验证结果: {self.validation_result}")
                    
                except json.JSONDecodeError:
                    logger.error(f"参数验证结果解析失败: {observation}")
            
            elif tool_name == "update_parameters":
                logger.info(f"参数更新结果: {observation}")
                # 参数更新已经在工具内部完成，这里只记录日志
    
    def _update_stage(self):
        """更新当前阶段"""
        # 如果会话已完成，不再更新阶段
        if self.is_session_completed:
            self.current_stage = "completed"
            return
            
        missing_params = self._get_missing_parameters()
        
        if missing_params:
            self.current_stage = "collection"
        elif not self.validation_result:
            self.current_stage = "validation"
        elif not self.validation_result.get("all_valid", False):
            self.current_stage = "collection"  # 返回收集阶段修正错误
        elif self._validation_outdated():
            # 如果验证结果过时（参数已更新但验证结果未同步），重新验证
            self.current_stage = "validation"
        else:
            self.current_stage = "confirmation"
    
    def _validation_outdated(self) -> bool:
        """检查验证结果是否过时（参数已更新但验证结果未同步）"""
        if not self.validation_result:
            return True
            
        validated_params = self.validation_result.get("validated_parameters", {})
        
        # 检查每个已收集的参数是否与验证结果中的参数一致
        for key, value in self.collected_parameters.items():
            if value is not None:
                if key not in validated_params or validated_params[key] != value:
                    logger.info(f"参数 {key} 已更新，需要重新验证: {validated_params.get(key)} -> {value}")
                    return True
        
        return False
    
    def _check_completion(self) -> bool:
        """检查是否完成参数收集"""
        missing_params = self._get_missing_parameters()
        
        return (
            len(missing_params) == 0 and
            self.validation_result is not None and
            self.validation_result.get("all_valid", False)
        )
    
    def _calculate_confidence_score(self) -> float:
        """计算置信度分数"""
        required_params = [
            'reservoir_name', 'forecast_model', 'start_time',
            'forecast_duration', 'rainfall_forecast_period', 'target_variable'
        ]
        
        # 参数完整性分数
        collected_count = sum(1 for param in required_params 
                            if param in self.collected_parameters and 
                            self.collected_parameters[param] is not None)
        completeness_score = collected_count / len(required_params)
        
        # 验证分数
        validation_score = 0.0
        if self.validation_result:
            validation_score = self.validation_result.get("validation_score", 0.0)
        
        # 综合分数
        return (completeness_score * 0.6 + validation_score * 0.4)
    
    def mark_session_completed(self):
        """标记会话完成"""
        self.is_session_completed = True
        self.current_stage = "completed"
        logger.info("会话已标记为完成")
    
    def reset_session(self):
        """重置会话状态"""
        self.collected_parameters = {}
        self.chat_history = []
        self.current_stage = "collection"
        self.validation_result = None
        self.is_session_completed = False
        self.agent_executor = None
        
        # 清理记忆功能
        self.memory.clear()
        
        logger.info("会话状态已重置")
    
    def get_session_summary(self) -> Dict[str, Any]:
        """获取会话摘要"""
        return {
            "collected_parameters": self.collected_parameters,
            "missing_parameters": self._get_missing_parameters(),
            "current_stage": self.current_stage,
            "validation_result": self.validation_result,
            "is_complete": self._check_completion(),
            "confidence_score": self._calculate_confidence_score(),
            "chat_history_length": len(self.chat_history),
            "timestamp": datetime.now().isoformat()
        }

# 使用示例
def main():
    """主函数，用于测试代理功能"""
    # 示例配置已移除OpenAI，只使用Ollama
    # openai_config = {
    #     "type": "openai",
    #     "model": "gpt-3.5-turbo",
    #     "temperature": 0.1,
    #     "api_key": "your-api-key",
    #     "base_url": "https://api.openai.com/v1"
    # }
    
    # 使用Ollama配置
    ollama_config = {
        "type": "ollama",
        "model": OLLAMA_CONFIG.get("chat_model", "llama2"),
        "temperature": 0.1,
        "base_url": OLLAMA_CONFIG.get("base_url", "http://10.48.0.81:11434")
    }
    
    # 创建代理实例
    agent = FloodForecastAgent(ollama_config)
    
    print("使用Ollama模型创建洪水预报代理...")
    print(f"模型: {ollama_config['model']}")
    print(f"服务地址: {ollama_config['base_url']}")
    
    # 测试对话
    test_inputs = [
        "我需要对大伙房水库进行洪水预报",
        "使用新安江模型，从明天上午8点开始",
        "预报48小时，降雨预见期12小时，目标变量是流量"
    ]
    
    for user_input in test_inputs:
        print(f"\n用户: {user_input}")
        response = agent.process_user_input(user_input)
        print(f"助手: {response['response']}")
        print(f"当前阶段: {response['current_stage']}")
        print(f"已收集参数: {response['collected_parameters']}")
        print(f"完成状态: {response['is_complete']}")


if __name__ == "__main__":
    main()